#! /usr/bin/python3
# -*- coding: utf-8 -*-
"""
Created on 2020/1/9 15:56 星期四

@author: jyz
"""
"""理论上实现每隔3s判断一次摄像头的设置是否有更新，若更新则kill所有已启动的进程并重启所有设备的所有算法"""

import multiprocessing as mp
import time
import os
import signal
import copy
import json
import cv2
from config.config_setting import ALL_PIDS
from config.getConfigData import getConfig
from config.config_setting import DEVICE_PIDS, logger, SCREENSHOT_IDX, SAVE_TIMESTAP, STREAM_FPS, STREAM_SIZE
from model import model_comm_func as cd
from utils import pipe_initial


class middleStation:

    # 初始化各种类
    def __init__(self):
        self.config_data = getConfig()
        self.cd = cd

    # 得到每个摄像头需要启动的算法
    def get_device_algorithm_config(self):
        camera_functions, dev = self.config_data.get_camera_functions()
        return camera_functions, dev

    def get_rtsp_rtmp(self):
        rtsp_paths, rtmp_paths = self.config_data.get_camera_url()
        return rtsp_paths, rtmp_paths

    def queue_img_put(self, input_quene, req, device_ids, discard_queue):  # 实现一次读取摄像头分给多个算法计算
        rtsp_paths, rtmp_paths = self.get_rtsp_rtmp()
        #print(rtsp_paths, rtmp_paths)

        caps = [cv2.VideoCapture(rtsp_paths[devid][0], cv2.CAP_FFMPEG) for devid in device_ids]
        [cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) for cap in caps]
        #print(caps)

        frameids = [0 for _ in range(len(device_ids))]
        pickframes = [0 for _ in range(len(device_ids))]
        fialedcap = set()
        while True:
            for i, devid in enumerate(device_ids):
                if not caps[i].isOpened():
                    logger.error('摄像头读取失败，重新获取摄像头...')
                    time.sleep(1)
                    video_path = rtsp_paths[devid][0]
                    print(video_path)
                    caps[i] = cv2.VideoCapture(video_path, cv2.CAP_FFMPEG)
                    print("111")
                    caps[i].set(cv2.CAP_PROP_BUFFERSIZE, 1)
                    print(caps[i])
                    continue
                
                ret = caps[i].grab()
                if not ret:
                    if devid not in fialedcap:
                        # fialedcap.add(devid)
                        logger.info(str(devid) + " isOpened but can not grab")
                    video_path = rtsp_paths[devid][0]
                    caps[i].release()
                    caps[i] = cv2.VideoCapture(video_path, cv2.CAP_FFMPEG)
                    continue
                frameids[i] += 1
                pickframe = int(frameids[i] * STREAM_FPS // 25)  # 设置输出帧数为STREAM_FPS，取倍数和余数
                if pickframe > pickframes[i]:
                    pickframes[i] = pickframe
                    ret, frame = caps[i].retrieve()
                    if not ret:
                        logger.info("XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX")
                    frame = cv2.resize(frame, STREAM_SIZE)
                    input_quene[-1][devid].put(frame)

                    for alg_type in req[devid].keys():  # 一次读取摄像头分给多个算法计算，防止多次读取rtsp，导致带宽不够。

                        if input_quene[alg_type][devid].qsize()<1:
                            input_quene[alg_type][devid].put(frame)
                        else:
                            try:
                                discard_queue[alg_type][devid].put(1, False)
                            except:
                                pass

                else:
                    pass
                if pickframes[i] == STREAM_FPS:
                    pickframes[i] = 0
                    frameids[i] = 0

    # 收集算法计算好的结果，用cv2画框，实现pipe推流和保存图片
    def push_orig_stream_2(self, q_push, device_):  # params: queue, pps
        c = 0
        rtsp_paths, rtmp_paths = self.get_rtsp_rtmp()

        pipes = {}
        for devid in device_:
            r = rtsp_paths[devid][1]
            if r: pipes[devid] = pipe_initial.PipeInitial(r)

        while True:
            for devid in device_:
                if not pipes.get(devid): continue
                if q_push[devid].qsize() < 1:
                    time.sleep(0.01)
                    continue
                img = q_push[devid].get()
                c += 1

                if pipes[devid].p.poll() is not None:  # poll的返回： A None value indicates that the process hasn’t terminated yet.
                    logger.info(str(pipes[devid].p.poll()) + 'ttt' * 8 + str(pipes[devid].p.pid))
                    try:
                        os.kill(pipes[devid].p.pid, signal.SIGKILL)
                    except:
                        pass
                    r = rtsp_paths[devid][1]
                    pipes[devid] = pipe_initial.PipeInitial(r)
                    logger.info("new pipe pid {}".format(pipes[devid].p.pid) * 10)
                if c % 1000 == 0:
                    logger.info('ccc'*10 + "{}".format(devid) + '+++'*10 + '{}'.format(pipes[devid].p.pid))
                    c = 0
                try:
                    for _ in range(1):
                        pipes[devid].p.stdin.write(img.tostring())
                except:
                    pass

    # 收集算法计算好的结果，实现pipe推流和保存图片
    def push_stream(self, q_push, algid, device_, discard_queue):  # params: queue, pps
        c = 0
        rtsp_paths, rtmp_paths = self.get_rtsp_rtmp()

        pipes = {}
        for devid in device_:
            r = rtmp_paths[devid].get(algid)
            if r: pipes[devid] = pipe_initial.PipeInitial(r)

        while True:
            for devid in device_:
                if not pipes.get(devid): continue
                if q_push[devid].qsize() < 1:
                    time.sleep(0.01)
                    continue
                img = q_push[devid].get()

                c += 1

                if pipes[devid].p.poll() is not None: # poll的返回： A None value indicates that the process hasn’t terminated yet.
                    logger.info(str(pipes[devid].p.poll()) + 'ttt' * 8 + str(pipes[devid].p.pid))
                    try:
                        os.kill(pipes[devid].p.pid, signal.SIGKILL)
                    except:
                        pass
                    r = rtmp_paths[devid][algid]
                    pipes[devid] = pipe_initial.PipeInitial(r)
                    logger.info("new pipe pid {}".format(pipes[devid].p.pid) * 10)
                if c % 1000 == 0:
                    logger.info('ccc'*10 + "{}".format(devid) + '+++'*10 + '{}'.format(pipes[devid].p.pid))
                    c = 0
                try:
                    for _ in range(1):
                        pipes[devid].p.stdin.write(img.tostring())
                    while discard_queue[algid][devid].qsize() > 0:
                        discard_queue[algid][devid].get()
                        pipes[devid].p.stdin.write(img.tostring())
                except:
                    pass

    def start_algorithm_queue(self, req, devices):  # device_config, device_ids; algos  [[10], [20]]
        self.init_save(devices)

        process_intrusion = []

        input_quene = {}
        predict_ans_quene = {}
        push_quene = {}
        discard_queue = {}
        input_quene[-1] = {}  # orig input

        for device in devices:
            input_quene[-1][device] = mp.Queue(maxsize=10)
        all_algorithm_types = []
        for dev_req in req:
            all_algorithm_types.extend(list(req[dev_req].keys()))

        all_algorithm_types = set(all_algorithm_types)
        for algo in all_algorithm_types:
            for que in [input_quene, predict_ans_quene, push_quene, discard_queue]:
                que[algo] = {}
                for device in devices:
                    que[algo][device] = mp.Queue(maxsize=10)
                que[algo][-1] = mp.Queue(maxsize=10)
        #### end of prepare quene

        # # 为每个算法创建输入队列
        process_intrusion.append(mp.Process(target=self.queue_img_put, args=(input_quene, req, devices, discard_queue)))  # why need coma ???

        for algo in all_algorithm_types:   # 一个算法需要做3件事，模型推理+判断推理结果+推流
            process_intrusion.append(mp.Process(target=self.cd.change_device_by_model,
                                                args=(input_quene[algo], predict_ans_quene[algo], algo, req)))  # origin_img_q, result_img_q, algo_id, device_config
            process_intrusion.append(mp.Process(target=self.cd.queue_img_get,
                                                args=(predict_ans_quene[algo], push_quene[algo], algo, req)))  # origin_img_q, result_img_q,algo_id, algo_params(timer)
            process_intrusion.append(mp.Process(target=self.push_stream, args=(push_quene[algo], algo, devices, discard_queue)))  # push_queue, algo_id, params, discard_queue
        process_intrusion.append(mp.Process(target=self.push_orig_stream_2, args=(input_quene[-1], devices)))  # push original stream

        [process.start() for process in process_intrusion]
        # 获取每个设备下的所有进程id，并添加到list中，方便后续重启时kill
        proce_ids = [process.pid for process in process_intrusion]
        ALL_PIDS.extend(proce_ids)
        print("ALL_PIDS {}".format(str(ALL_PIDS)))

    def init_save(self, devices):
        for device in devices:
            SCREENSHOT_IDX.setdefault(device, {})
            SAVE_TIMESTAP.setdefault(device, {})

    def restart_device(self, cameras_function, device):
        rtsp_paths, rtmp_paths = self.get_rtsp_rtmp()

        # 先用杀掉该设备之前启动的进程
        self.cd.kill_started_pro(device)

        [rtsp_path, src_push_url], rtmp_path = rtsp_paths[device], rtmp_paths[device]
        logger.info("{}正在重启的设备rtmp地址{}".format(cameras_function, rtmp_path))

        # self.start_algorithm_queue(cameras_function[device], rtsp_path, rtmp_path, src_push_url, device)

    def update_device_config_2(self):
        pre_device_alg = None
        while True:
            cur_config, devi = self.get_device_algorithm_config()

            if pre_device_alg is None:
                pre_device_alg = copy.deepcopy(cur_config)
            elif pre_device_alg == cur_config:
                pass
            else:
                pre_device_alg = copy.deepcopy(cur_config)
                logger.info('配置改变！')
                print("ALL_PIDS {}".format(str(ALL_PIDS)))
                for pid in ALL_PIDS:
                    try:
                        os.kill(pid, signal.SIGKILL)
                    except:
                        pass
                ALL_PIDS.clear()
                # =========== restart ==============
                devices_config, devices = self.get_device_algorithm_config()
                print(devices_config)
                self.start_algorithm_queue(devices_config, devices)
                # =========== restart ==============
            time.sleep(3)  # 每隔3秒请求一次，对比设备的参数。


if __name__ == '__main__':

    am = middleStation()
    devices_config, devices = am.get_device_algorithm_config()
    am.start_algorithm_queue(devices_config, devices)
    am.update_device_config_2()

